import asyncio
import os
import streamlit as st
from agents.agent import Agent
from agents.mcp_client import MCPClient

st.title("MCP Agent 调试")
agent = None
# agent = BaseAgent(name="MCP Agent", description="一个编程智能助手")
# agent = Agent()
async def load_agent_tools():
    
    await agent._load_tools()

def init_agent():
        # 初始化session_state
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "agent" not in st.session_state:
        st.session_state.agent = agent
        
def chat_with_agent(prompt: str):
    if prompt:
        # 显示用户消息
        st.chat_message("user").markdown(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # 如果agent已初始化，调用agent
        if st.session_state.agent:
            try:
                with st.chat_message("assistant"):
                    with st.spinner("思考中..."):
                        response = st.session_state.agent.call_llm(prompt)
                    st.markdown(response)
                st.session_state.messages.append({"role": "assistant", "content": response})
            except Exception as e:
                st.error(f"调用Agent时出错: {str(e)}")
        else:
            st.warning("请先配置API密钥以初始化Agent")

# 创建一个异步函数来处理工具调用
async def handle_tool_call(input_text, tool_name):
    result = await agent.call_llm_with_mcp_tools(input_text, tool_name)
    return result

async def display_with_agent():
    
    enabled_mcp_clients = []
    # agent = None
    try:
        # for mcp_tool in [
        #     PresetMcpTools.filesystem.append_mcp_params(f" {PROJECT_ROOT_DIR!s}"),
        #     PresetMcpTools.fetch,
        # ]:
        #     rprint(mcp_tool.shell_cmd)
        #     mcp_client = MCPClient(**mcp_tool.to_common_params())
        #     enabled_mcp_clients.append(mcp_client)

        mcp_client = MCPClient(
            name="calculate",
            command="python",
            args=["./mcp/mcp_server.py"],
        )
        await mcp_client.init()
        enabled_mcp_clients.append(mcp_client)

        agent = Agent(
            model=os.getenv("LLM_MODEL_ID"),
            mcp_clients=enabled_mcp_clients,
        )
        await agent.init()
        
    except Exception as e:
        st.error(f"初始化Agent时出错: {str(e)}")
        
    with st.sidebar:
        st.header("设置")
    
        tools = st.selectbox(
            label="选择工具",
            # options=["calculate", "get_current_time", "auto_generate_word"],
            options=[tool.name for tool in agent.llm.tools],
            index=0
        )
        if tools == "calculate_num":
            expression = st.text_input("输入表达式", "1 + 2 * 3")
        elif tools == "get_current_time":
            # 使用 asyncio.run 来运行异步函数
            current_time = asyncio.run(handle_tool_call(input, tools))
            st.write(current_time)
        else:
            prompt = st.text_input("输入提示", "一个随机单词")
        
    # 用户输入
    init_agent()
    
    prompt = st.chat_input("请输入您的问题...")
    chat_with_agent(prompt)
    

# asyncio.run(load_agent_tools())
# display_with_agent()
    


try:
    asyncio.run(load_agent_tools())
except Exception as e:
    st.error(f"加载工具时出错: {str(e)}")
    
asyncio.run(display_with_agent())
